from matplotlib import rc
params = {'backend': 'AGG',
          'axes.labelsize': 10,'text.fontsize': 10,
          'legend.fontsize': 10,
          'xtick.labelsize': 8,
          'ytick.labelsize': 8,
          'text.usetex': True}
rc(params)

import matplotlib.pyplot as plt
from numpy import loadtxt,ones,amin,amax

class LinePlot:
    'Drawing class for this benchmark.'
    def __init__(self):
        self.extension = '.pdf'
        self.annotations = {0 : {'label' : r'$t$',               'title' : '', 'name' : ''},
                            1 : {'label' : r'$dt$',              'title' : 'Time Step', 'name' : 'dt'},
                            2 : {'label' : r'$V_{a}$',           'title' : 'Volume of Approx. Solution', 'name' : 'av'},
                            3 : {'label' : r'$V_{a}^{E},\,\%$','title' : 'Error of Volume of Approx. Solution', 'name' : 'ave'},
                            4 : {'label' : r'$V_{e}$',           'title' : 'Volume of Exact Solution', 'name' : 'ev'},
                            5 : {'label' : r'$V_{e}^{E},\,\%$','title' : 'Error of Volume of Exact Solution', 'name' : 'eve'},
                            6 : {'label' : r'$E$',               'title' : 'L2 Error', 'name' : 'e'},
                            7 : {'label' : r'$EE$',               'title' : 'Enright error', 'name' : 'enright'}
                            }

    def plot(self,x,y,Y,xlabel,ylabel,title,name):
        '''Simple time vs. quantity plot.'''
        fig = plt.figure()
        plt.plot(x,y,'b:o')
        plt.plot(x,Y*ones(len(x)),'r')
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        xmin = amin(x)
        xmax = amax(x)
        ymin = amin(y)
        ymax = amax(y)

        dy = abs(ymax-ymin)/2.
        if dy == 0:
            plt.axis('tight')
        else:
            plt.xlim([xmin,xmax])
            plt.ylim([ymin-dy,ymax+dy])

        plt.savefig(name,dpi=(300),bbox_inches='tight')

    def plot_all(self,problem,solver):
        '''Plots all active variables.'''
        saveDir = solver.saveDir
        transFile = '/'.join([saveDir,'_'.join([str(solver.CFL),problem.outFileT])])
        avgFile = '/'.join([saveDir,'_'.join([str(solver.CFL),problem.outFileA])])
        allCols = loadtxt(transFile)   #(rows,cols)
        allAvgs = loadtxt(avgFile)     #(cols,)

        if allCols.shape[1] != allAvgs.shape[0]:
            raise ValueError('#columns for averages(%s) and transients(%d) does not agree' % (allAvgs.shape[0],allCols.shape[1]))
        
        x = allCols[:,0]
        xlabel = self.annotations[0]['label']
        for i in problem.activeVariables[1:]:
            y = allCols[:,i]
            Y = allAvgs[i]
            ylabel = self.annotations[i]['label']
            title = self.annotations[i]['title']
            
            name = '_'.join([str(solver.CFL),problem.problemDir[0],solver.methodDir[0],solver.solverDir[0] + solver.solverDir[-1]])
            name = '_'.join([name,self.annotations[i]['name']+self.extension])
            name = '/'.join([saveDir,name])

            self.plot(x,y,Y,xlabel,ylabel,title,name)
